#!/bin/bash
# QEMU管理脚本
#
# 向后兼容性说明:
# - 原有核心命令 (run, copy, monitor) 保持完全兼容，支持多平台
# - 支持平台: m64x86, aarch64_ex, aarch64
# - 新增便利命令提供更细粒度的控制
#
# 使用示例:
#   ./qemu.sh run aarch64     # 运行aarch64平台 (兼容原有用法)
#   ./qemu.sh copy aarch64    # 复制文件到aarch64 (兼容原有用法)  
#   ./qemu.sh monitor         # 连接监控 (兼容原有用法)

QEMU_MONITOR="/tmp/qemu-monitor"

function qemu_create_usb() {
  local platform=$1
  local usb_disk_path="device/qemu/$platform/usb-disk.raw"
  
  if [ ! -f "$usb_disk_path" ]; then
    echo "Creating USB disk image at $usb_disk_path..."
    mkdir -p "device/qemu/$platform"
    # 创建一个 512MB 的虚拟 U 盘
    dd if=/dev/zero of="$usb_disk_path" bs=1M count=512
    # 格式化为 FAT32 文件系统
    mkfs.vfat "$usb_disk_path"
    echo "USB disk image created successfully"
  else
    echo "USB disk image already exists at $usb_disk_path"
  fi
}

function qemu_create_rootfs_from_tarball() {
  local platform=$1
  local rootfs_ext4_path="device/qemu/$platform/rootfs.ext4"
  local rootfs_tarball=""
  
  # 查找可能的 rootfs.tar.gz 文件位置
  local possible_locations=(
    "device/qemu/$platform/rootfs.tar.gz"
  )
  
  for location in "${possible_locations[@]}"; do
    if [ -f "$location" ]; then
      rootfs_tarball="$location"
      echo "Found rootfs tarball at: $rootfs_tarball"
      break
    fi
  done
  
  if [ -z "$rootfs_tarball" ]; then
    echo "Error: No rootfs.tar.gz found in any of the expected locations:"
    printf "  - %s\n" "${possible_locations[@]}"
    return 1
  fi
  
  if [ -f "$rootfs_ext4_path" ]; then
    echo "rootfs.ext4 already exists at $rootfs_ext4_path"
    read -p "Do you want to recreate it? (y/N): " answer
    if [[ ! "$answer" =~ ^[Yy]$ ]]; then
      echo "Skipping rootfs.ext4 creation"
      return 0
    fi
    rm -f "$rootfs_ext4_path"
  fi
  
  echo "Creating rootfs.ext4 from $rootfs_tarball..."
  mkdir -p "device/qemu/$platform"
  
  # 创建临时目录来解压 rootfs
  local temp_dir=$(mktemp -d)
  echo "Extracting rootfs to temporary directory: $temp_dir"
  
  # 解压 rootfs.tar.gz
  if ! tar -xzf "$rootfs_tarball" -C "$temp_dir"; then
    echo "Error: Failed to extract $rootfs_tarball"
    rm -rf "$temp_dir"
    return 1
  fi
  
  # 计算需要的镜像大小（比实际内容大 20% 作为缓冲）
  local content_size=$(du -sb "$temp_dir" | cut -f1)
  local image_size=$(( content_size * 120 / 100 ))
  local image_size_mb=$(( image_size / 1024 / 1024 + 100 )) # 至少增加100MB
  
  echo "Content size: ${content_size} bytes"
  echo "Creating ${image_size_mb}MB ext4 image..."
  
  # 创建空的镜像文件
  dd if=/dev/zero of="$rootfs_ext4_path" bs=1M count="$image_size_mb"
  
  # 格式化为 ext4 文件系统
  mkfs.ext4 -F "$rootfs_ext4_path"
  
  # 创建临时挂载点
  local mount_point=$(mktemp -d)
  echo "Mounting image at: $mount_point"
  
  # 挂载镜像
  if ! sudo mount -o loop "$rootfs_ext4_path" "$mount_point"; then
    echo "Error: Failed to mount $rootfs_ext4_path"
    rm -rf "$temp_dir" "$mount_point"
    return 1
  fi
  
  # 复制内容到镜像
  echo "Copying rootfs contents to ext4 image..."
  if ! sudo cp -a "$temp_dir"/* "$mount_point"/; then
    echo "Error: Failed to copy contents to image"
    sudo umount "$mount_point"
    rm -rf "$temp_dir" "$mount_point"
    return 1
  fi
  
  # 卸载并清理
  sudo umount "$mount_point"
  rm -rf "$temp_dir" "$mount_point"
  
  echo "Successfully created rootfs.ext4 at $rootfs_ext4_path"
  echo "Image size: ${image_size_mb}MB"
  
  # 检查生成的镜像
  echo "Verifying the created image..."
  if ! sudo e2fsck -n "$rootfs_ext4_path"; then
    echo "Warning: Image verification failed, but image was created"
  else
    echo "Image verification passed"
  fi
}

function qemu_plug_usb() {
  local platform=${1:-aarch64}
  local usb_disk_path="device/qemu/$platform/usb-disk.raw"
  
  # 确保 U 盘镜像存在
  qemu_create_usb "$platform"
  
  echo "drive_add 0 file=$usb_disk_path,if=none,id=usbdrive" | sudo socat - UNIX-CONNECT:${QEMU_MONITOR}
  echo "device_add usb-storage,bus=xhci.0,drive=usbdrive,id=usb1" | sudo socat - UNIX-CONNECT:${QEMU_MONITOR}
}

function qemu_unplug_usb() {
  echo "device_del usb1" | sudo socat - UNIX-CONNECT:${QEMU_MONITOR}
}

function qemu_quit() {
  echo "quit" | sudo socat - UNIX-CONNECT:${QEMU_MONITOR}
}

function qemu_copy() {
  PLATFORM=$1
  TARGET_IMAGE="device/qemu/$PLATFORM/rootfs.ext4"

  if [ ! -f "$TARGET_IMAGE" ]; then
    echo "Target image $TARGET_IMAGE does not exist, attempting to create from rootfs.tar.gz..."
    qemu_create_rootfs_from_tarball "$PLATFORM"
    if [ $? -ne 0 ]; then
      echo "Error: Failed to create rootfs.ext4"
      return 1
    fi
  fi

  # 检查镜像是否可以正常挂载
  echo "Checking if the image can be mounted..."
  if ! sudo e2fsck -n "$TARGET_IMAGE" &>/dev/null; then
    echo "Error: The image file system is corrupted or cannot be mounted"
    return 1
  fi

  echo "Checking if root directory is accessible..."
  if ! sudo e2ls "$TARGET_IMAGE:/" &>/dev/null; then
    echo "Error: Cannot access root directory of the image"
    return 1
  fi

  # 使用固定路径配置（后续添加新路径只需在此数组新增条目）
  # 注意：QEMU模拟单分区环境，实际硬件使用A/B双分区设计
  local directories_to_copy=(
    "./device/qemu/overlay/usr:/usr"               # 用户程序和配置
    "./device/qemu/overlay/mnt:/mnt"               # 挂载点(包含assets静态资源)
    "./device/qemu/overlay/etc:/etc"               # 用户配置文件(覆盖)
    "./out/web:/mnt/assets/web"                    # Web资源目录(构建输出)
    "./out/bin:/usr/bin"                           # 主程序二进制文件(构建输出)
    "./third_party/lib/aarch64:/usr/lib"           # aarch64动态库目录
  )

  # 统一处理所有需要复制的目录
  for entry in "${directories_to_copy[@]}"; do
    local src="${entry%%:*}"
    local dst="${entry##*:}"

    if [ -d "$src" ]; then
      copy_to_image "$src" "$dst" "$TARGET_IMAGE"
    else
      echo "Note: Source directory $src does not exist, skipping..."
    fi
  done

  echo "Checking and repairing file system after copying..."
  sudo e2fsck -fy "$TARGET_IMAGE"

  # 强制卸载文件系统
  echo "Forcing unmount of the file system..."
  sudo umount -f /tmp/mnt 2>/dev/null || true

  echo "File copying completed for $PLATFORM"
}

# 新增通用复制函数
function copy_to_image() {
  local src_dir="$1"
  local target_path="$2"
  local image="$3"

  echo "Starting to copy $src_dir to $image:$target_path"

  # 删除旧目录时需要逐级处理
  echo "Deleting old $target_path directory..."
  sudo e2rm -rf "$image:$target_path" 2>/dev/null || true

  # 修改目录创建方式（使用手动逐级创建）
  local current_path=""
  IFS='/' read -ra ADDR <<<"$target_path"
  for i in "${ADDR[@]}"; do
    [ -z "$i" ] && continue
    current_path="$current_path/$i"
    if ! sudo e2mkdir "$image:$current_path" 2>/dev/null; then
      echo "Note: Directory $current_path already exists"
    fi
  done

  # 执行递归复制
  copy_recursive "$src_dir" "$target_path" "$image"
}

# 修改后的递归复制函数（保持文件权限）
function copy_recursive() {
  local src="$1"
  local dst="$2"
  local image="$3"

  for item in "$src"/*; do
    local name=$(basename "$item")
    if [ -d "$item" ]; then
      # 仅创建当前层级目录
      if ! sudo e2mkdir "$image:$dst/$name" 2>/dev/null; then
        echo "Note: Directory $dst/$name already exists"
      fi
      copy_recursive "$item" "$dst/$name" "$image"
    elif [ -f "$item" ]; then
      # 复制文件并保持权限
      if ! sudo e2cp -p "$item" "$image:$dst/$name"; then
        echo "Error: Failed to copy $item to $dst/$name"
      else
        echo "Copied: $item to $dst/$name"
        # 如果是shell脚本或其他可执行文件，确保设置执行权限
        if [[ "$name" == *.sh ]] || [[ "$name" == S[0-9][0-9]* ]] || [[ -x "$item" ]]; then
          echo "Setting executable permission for $dst/$name"
          # 注意：e2cp的-p选项应该保持权限，但为了确保执行权限，这里可以添加额外检查
        fi
      fi
    fi
  done
}

function qemu_run() {
  PLATFORM=$1
  case $PLATFORM in
  m64x86)
    sudo qemu-system-x86_64 -M pc -kernel device/qemu/m64x86/bzImage \
      -drive file=device/qemu/m64x86/rootfs.ext2,if=virtio,format=raw \
      -append "rootwait root=/dev/vda console=tty1 console=ttyS0" \
      -serial stdio -net nic,model=virtio -net user
    ;;
  aarch64_ex)
    sudo qemu-system-aarch64 -M virt -cpu cortex-a53 -nographic -smp 2 -m 2048 \
      -kernel device/qemu/aarch64/Image \
      -append "rootwait root=/dev/vda console=ttyAMA0" \
      -netdev user,id=eth0 -device virtio-net-device,netdev=eth0 \
      -drive file=device/qemu/aarch64/rootfs.ext4,if=none,format=raw,id=hd0 \
      -device virtio-blk-device,drive=hd0 \
      -device qemu-xhci,id=xhci \
      -monitor unix:/tmp/qemu-monitor,server,nowait \
      -no-shutdown \
      -serial mon:stdio
    ;;
  aarch64)
    # 配置网桥环境
    qemu_setup_bridge ens37
    
    # 检查 rootfs.ext4 是否存在，如果不存在则从 tar.gz 创建
    if [ ! -f "device/qemu/aarch64/rootfs.ext4" ]; then
      echo "rootfs.ext4 not found, attempting to create from rootfs.tar.gz..."
      qemu_create_rootfs_from_tarball aarch64
      if [ $? -ne 0 ]; then
        echo "Error: Failed to create rootfs.ext4"
        exit 1
      fi
    fi
    
    # 确保 U 盘镜像存在
    qemu_create_usb aarch64
    
    sudo qemu-system-aarch64 -M virt -cpu cortex-a53 -nographic -smp 2 -m 2048 \
      -kernel device/qemu/aarch64/Image \
      -append "rootwait root=/dev/vda console=ttyAMA0" \
      -netdev bridge,id=eth0,br=br0 -device virtio-net-device,netdev=eth0 \
      -drive file=device/qemu/aarch64/rootfs.ext4,if=none,format=raw,id=hd0 \
      -device virtio-blk-device,drive=hd0 \
      -drive file=device/qemu/aarch64/usb-disk.raw,if=none,format=raw,id=usbdrive \
      -device qemu-xhci,id=xhci \
      -device usb-storage,bus=xhci.0,drive=usbdrive,id=usb1 \
      -monitor unix:/tmp/qemu-monitor,server,nowait \
      -no-shutdown \
      -serial mon:stdio
    ;;
  *)
    echo "Error: Unsupported platform for QEMU: $PLATFORM"
    exit 1
    ;;
  esac
}

function qemu_monitor() {
  sudo socat - UNIX-CONNECT:/tmp/qemu-monitor
}

# Ubuntu 24 网桥配置函数 - 桥接到物理网卡
function qemu_setup_bridge() {
  local interface_name="${1:-ens37}"  # 默认使用ens37，可通过参数自定义
  
  echo "Setting up br0 bridge bridged to $interface_name for Ubuntu 24..."
  
  if ! ip link show br0 &>/dev/null; then
    # 检查指定网卡是否存在
    if ! ip link show $interface_name &>/dev/null; then
      echo "Error: $interface_name interface not found!"
      echo "Available interfaces:"
      ip link show | grep -E "^[0-9]+:" | cut -d: -f2 | tr -d ' '
      return 1
    fi
    
    # 保存网卡当前的IP配置
    local INTERFACE_IP=$(ip addr show $interface_name | grep -oP 'inet \K[\d.]+/\d+' | head -1)
    local INTERFACE_GW=$(ip route | grep "default.*$interface_name" | awk '{print $3}' | head -1)
    
    echo "Current $interface_name IP: $INTERFACE_IP, Gateway: $INTERFACE_GW"
    
    # 创建网桥
    sudo ip link add name br0 type bridge
    sudo ip link set dev br0 up
    
    # 将网卡加入网桥
    echo "Adding $interface_name to bridge br0..."
    sudo ip link set dev $interface_name master br0
    
    # 如果网卡有IP地址，将其移动到网桥上
    if [ -n "$INTERFACE_IP" ]; then
      echo "Moving IP $INTERFACE_IP from $interface_name to br0..."
      sudo ip addr del $INTERFACE_IP dev $interface_name 2>/dev/null || true
      sudo ip addr add $INTERFACE_IP dev br0
      
      # 删除旧的默认路由再添加新的，避免重复路由
      if [ -n "$INTERFACE_GW" ]; then
        echo "Removing old default routes..."
        sudo ip route del default 2>/dev/null || true
        echo "Adding default gateway via br0..."
        sudo ip route add default via $INTERFACE_GW dev br0
      fi
    fi
    
    # 启用IP转发
    echo 1 | sudo tee /proc/sys/net/ipv4/ip_forward > /dev/null
    
    echo "Created br0 bridge with $interface_name attached"
    echo "QEMU guests will get IP from external DHCP server"
  else
    echo "br0 bridge already exists"
    # 确保网桥是启用状态
    sudo ip link set dev br0 up
    
    # 检查指定网卡是否存在
    if ip link show $interface_name &>/dev/null; then
      # 检查网卡是否已经在网桥中
      if ! ip link show $interface_name | grep -q "master br0"; then
        echo "Adding $interface_name to existing br0 bridge..."
        sudo ip link set dev $interface_name master br0
        
        # 检查网桥是否有IP地址，如果没有则从网卡移动过来
        local BRIDGE_IP=$(ip addr show br0 | grep -oP 'inet \K[\d.]+/\d+' | head -1)
        if [ -z "$BRIDGE_IP" ]; then
          local INTERFACE_IP=$(ip addr show $interface_name | grep -oP 'inet \K[\d.]+/\d+' | head -1)
          local INTERFACE_GW=$(ip route | grep "default.*$interface_name" | awk '{print $3}' | head -1)
          
          if [ -n "$INTERFACE_IP" ]; then
            echo "Moving IP $INTERFACE_IP from $interface_name to br0..."
            sudo ip addr del $INTERFACE_IP dev $interface_name 2>/dev/null || true
            sudo ip addr add $INTERFACE_IP dev br0
            
            # 更新路由
            if [ -n "$INTERFACE_GW" ]; then
              echo "Updating default gateway via br0..."
              sudo ip route del default dev $interface_name 2>/dev/null || true
              sudo ip route add default via $INTERFACE_GW dev br0 2>/dev/null || true
            fi
          fi
        fi
      else
        echo "$interface_name already bridged to br0"
      fi
    else
      echo "Note: $interface_name interface not found, skipping bridge attachment"
      echo "Bridge br0 is ready for QEMU use"
    fi
  fi
}

# 新增：网桥清理函数
function qemu_cleanup_bridge() {
  local interface_name="${1:-ens37}"
  
  echo "Cleaning up br0 bridge..."
  
  if ip link show br0 &>/dev/null; then
    # 保存网桥的IP配置
    local BRIDGE_IP=$(ip addr show br0 | grep -oP 'inet \K[\d.]+/\d+' | head -1)
    local BRIDGE_GW=$(ip route | grep "default.*br0" | awk '{print $3}' | head -1)
    
    # 检查网卡是否在网桥中
    if ip link show $interface_name &>/dev/null && ip link show $interface_name | grep -q "master br0"; then
      echo "Removing $interface_name from bridge br0..."
      sudo ip link set dev $interface_name nomaster
      
      # 将IP地址移回原网卡
      if [ -n "$BRIDGE_IP" ]; then
        echo "Moving IP $BRIDGE_IP from br0 to $interface_name..."
        sudo ip addr del $BRIDGE_IP dev br0 2>/dev/null || true
        sudo ip addr add $BRIDGE_IP dev $interface_name
        
        # 恢复默认路由
        if [ -n "$BRIDGE_GW" ]; then
          echo "Restoring default gateway via $interface_name..."
          sudo ip route del default 2>/dev/null || true
          sudo ip route add default via $BRIDGE_GW dev $interface_name
        fi
      fi
    fi
    
    # 删除网桥
    sudo ip link set dev br0 down
    sudo ip link delete br0
    echo "Bridge br0 removed successfully"
  else
    echo "Bridge br0 does not exist"
  fi
}